#define _CRT_SECURE_NO_WARNINGS
#define _WINSOCK_DEPRECATED_NO_WARNINGS
#pragma comment (lib, "ws2_32.lib")

#include <stdio.h>
#include <winsock2.h>
#include <time.h>

#include "DNS.h"

#define DNS_PORT 53 // DNS的端口号，这是固定的

void print_info();                                                                                                                                                                                               //输出本次实验相关信息
int process_parameters(int argc, char* argv[], char* DNS_Server_IP, char** file_path);                                                                                                                                                                //解析程序运行方式
void Receive_from_Local(SOCKET local_sock, SOCKET extern_sock, sockaddr_in extern_Addr, Cache cache, Local_Table local_table, Cache_LRU_List cache_lru_list, ID_Trans_Table** id_trans_table, int debug_level);  //接受本地查询请求
void Receive_from_Extern(SOCKET local_sock, SOCKET extern_sock, sockaddr_in extern_Addr, Cache cache, Local_Table local_table, Cache_LRU_List cache_lru_list, ID_Trans_Table** id_trans_table, int debug_level); //接受向服务器查询的结果
int main(int argc, char *argv[])
{
    //初始化 DLL,主版本与副版本都选2
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);
    char DNS_Server_IP[16] = "192.168.1.1";
    Local_Table local_table=create_hash_map();
    Cache cache=create_hash_map();
    Cache_LRU_List cache_lru_list=create_list();
    //用数组记录向DNS服务器查询的记录,数组下标+1即其向服务器查询的ID
    ID_Trans_Table **id_trans_table = (ID_Trans_Table **)malloc(sizeof(ID_Trans_Table *) * MAX_ID_TRANS_TABLE_SIZE);
    for (int i = 0; i < MAX_ID_TRANS_TABLE_SIZE; i++)
        id_trans_table[i] = NULL;
    print_info();
    char* file_path=NULL;
    int debug_level = process_parameters(argc, argv,DNS_Server_IP,&file_path);
    //创建套接字,选择UDP协议（DGRAM与STREAM分别对应UDP与TCP,协议为0时根据传输方式自动选择协议）
    SOCKET local_sock = socket(AF_INET, SOCK_DGRAM, 0), extern_sock = socket(AF_INET, SOCK_DGRAM, 0);
    //绑定套接字（extern_Addr_in相当于是sockadd的针对IPv4特例,sockadd的data部分包含地址等信息但难以赋值操作,故先创建sockaddr_in,再将其强制转化）
    //分别为中继器地址与DNS服务器地址
    sockaddr_in extern_Addr, local_Addr;
    memset(&extern_Addr, 0, sizeof(extern_Addr));           //每个字节都用0填充
    extern_Addr.sin_family = AF_INET;                       //使用IPv4地址
    extern_Addr.sin_addr.s_addr = inet_addr(DNS_Server_IP); //具体的IP地址
    extern_Addr.sin_port = htons(DNS_PORT);                 //端口号
    memset(&local_Addr, 0, sizeof(local_Addr));
    local_Addr.sin_family = AF_INET;
    local_Addr.sin_addr.s_addr = INADDR_ANY;
    local_Addr.sin_port = htons(DNS_PORT);
     //设置一对多
    int reuse = 1;
    setsockopt(local_sock, SOL_SOCKET, SO_REUSEADDR, (const char*)&reuse, sizeof(reuse));
    bind(local_sock, (struct sockaddr*)&local_Addr, sizeof(local_Addr));
    if (debug_level>=1)
        printf("bind successfully.\n");
    int read_result = read_from_txt(local_table, debug_level,file_path);
    if (debug_level >= 1 && read_result != -1)
        printf("read from txt successfully.\n");
    else if (read_result == -1)
        printf("read error or read nothing");
    while (true)
    {
        fd_set fread = { 0 };
        timeval timeout = { 0 };
        FD_ZERO(&fread);
        FD_SET(extern_sock, &fread);
        FD_SET(local_sock, &fread);
        timeout.tv_usec = 10000;
        int nRet2 = select(0, &fread, NULL, NULL, &timeout);
        if (FD_ISSET(extern_sock, &fread))
            Receive_from_Extern(local_sock, extern_sock, extern_Addr, cache, local_table, cache_lru_list, id_trans_table, debug_level);
        if (FD_ISSET(local_sock, &fread))
            Receive_from_Local(local_sock, extern_sock, extern_Addr, cache, local_table, cache_lru_list, id_trans_table, debug_level);
    }
}

void print_info()
{
    printf("************************************************************************\n");
    printf("*           @Course Name:Course Design of Computer Network             *\n");
    printf("*           @Name of Team members:luo yi,zhang lei,zhang youkang       *\n");
    printf("*           @Teacher:gao zhanchun      @Class number:2020211307        *\n");
    printf("* -------------------------------------------------------------------- *\n");
    printf("*                        DNS Relay Server - Ver 1.0                    *\n");
    printf("************************************************************************\n");
    printf("Command syntax:dnsrelay [-d | -dd] [dns-server-IP-addr] [URL-IP_table]  \n");
    return;
}

int process_parameters(int argc, char *argv[], char *DNS_Server_IP,char **file_path)
{
    int debug_level=0;
    int user_set_dns_flag = 0;
    int user_set_file_path = 0;
    if (argc > 1 && argv[1][0] == '-')
    {
        if (argv[1][1] == 'd')
            debug_level++; 
        if (argv[1][2] == 'd')
            debug_level++; 
        if (argc > 2)
        {
            user_set_dns_flag = 1; 
            strcpy(DNS_Server_IP, argv[2]);
        }
        if (argc > 3)
        {
            user_set_file_path = 1;
            *file_path = argv[3];
        }
    }
    if (user_set_dns_flag) 
        printf("Set DNS server:%s\n", argv[2]);
    else
        printf("Set DNS server:%s by default\n", DNS_Server_IP);
    if (user_set_file_path)
        printf("Set URL-IP table path:%s\n", argv[3]);
    else
        printf("Set URL-IP table path by default\n");
    printf("Debug level:%d\n", debug_level);
    return debug_level;
}
void Receive_from_Local(SOCKET local_sock, SOCKET extern_sock, sockaddr_in extern_Addr, Cache cache, Local_Table local_table, Cache_LRU_List cache_lru_list, ID_Trans_Table **id_trans_table, int debug_level)
{
    char buf[BUF_MAX_SIZE];
    char ip[16]="", url[65]="";
    sockaddr_in client;
    int length = -1, length_client = sizeof(client);
    memset(buf, 0, BUF_MAX_SIZE);
    int ipv6_flag = 0;
    //将数据接收到缓冲区中,套接口为SOCK_DGRAM类型,则发送数据源的地址被复制到相应的sockaddr结构中
    length = recvfrom(local_sock, buf, sizeof buf, 0, (struct sockaddr *)&client, &length_client);
    if (length > 0)
    {
        //判断IPV6
        if (get_url_from_question(buf, url, debug_level) == -1) {
            ipv6_flag = 1;
        }

        get_url_from_question(buf,url,debug_level);
        if (debug_level >= 1)
        {
            printf("\n\n---- Recv:Client [IP:%s]----\n", inet_ntoa(client.sin_addr));
            time_t t = time(NULL);
            char temp[64];
            strftime(temp, sizeof(temp), "%Y/%m/%d %X %A", localtime(&t));
            printf("%s\n", temp);
            printf("Receive from client [Query:%s]\n", url);
        }
        /* 依次在cache、内存中查找,最后询问服务器 */
        char* value = hash_map_get(cache, url);
        if (value!=NULL&&ipv6_flag!=1)
        {
            strcpy(ip,value);
            list_remove(cache_lru_list, url);
            list_pushback(cache_lru_list, url);
            if (debug_level >= 1)
            {
                printf("Read from cache [Url:%s -> IP:%s]\n", url, ip);
                printf_IPmap(cache);
            }
        }
        else
        {
            value = hash_map_get(local_table, url);
            if ((value != NULL && ipv6_flag != 1)||((value != NULL && (strcmp(value, "0.0.0.0") == 0) && ipv6_flag == 1)))
            {
                strcpy(ip, value);
                if (debug_level >= 1)
                    printf("Read from local data [Url:%s -> IP:%s]\n", url, ip);
            }
            else
            {
                printf("[Url:%s] not in local data and cache\n", url);
                unsigned short index = 0;
                while (index < MAX_ID_TRANS_TABLE_SIZE && id_trans_table[index] != NULL)
                {
                    if (Check_ID_Expired(id_trans_table[index]))
                    {
                        ID_Trans_Table *current = id_trans_table[index];
                        free(current);
                        id_trans_table[index] = NULL;
                        break;
                    }
                    index++;
                }
                if (index >= MAX_ID_TRANS_TABLE_SIZE)
                {
                    if (debug_level)
                        printf("Register failed, the ID transfer table is full.\n");
                    return;
                }

                id_trans_table[index] = (ID_Trans_Table *)malloc(sizeof(ID_Trans_Table));
                memcpy(&(id_trans_table[index]->old_ID), buf, (size_t)2);//保存请求报文id
                Set_ID_Expire(id_trans_table[index], ID_EXPIRE_TIME);
                id_trans_table[index]->client = client;
                index++;
                memcpy(buf, &index, (size_t)2);//替换请求报文id并转发
                length = sendto(extern_sock, buf, length, 0, (struct sockaddr *)&extern_Addr, sizeof(extern_Addr));
                if (debug_level >= 1)
                    printf("Send to external DNS server [Url:%s]\n", url);
                return;
            }
        }
        //不需要中转则直接构造DNS报文
        char sendbuf[BUF_MAX_SIZE];
        int send_len = write_package(sendbuf, buf, length, ip, debug_level);
        int real_send_len = sendto(local_sock, sendbuf, send_len, 0, (SOCKADDR *)&client, sizeof(client));
        if (real_send_len < send_len)
            printf("Error:Send packet to client");
        if (debug_level >= 1)
            printf("Send packet to client[Url:%s -> IP:%s]\n", url, ip);
    }
    return;
}
void Receive_from_Extern(SOCKET local_sock, SOCKET extern_sock, sockaddr_in extern_Addr, Cache cache, Local_Table local_table, Cache_LRU_List cache_lru_list, ID_Trans_Table **id_trans_table, int debug_level)
{
    char buf[BUF_MAX_SIZE];
    sockaddr_in external, client;
    memset(buf, 0, BUF_MAX_SIZE);
    int length = -1, length_extern = sizeof(external);
    char ip[16]="", url[65]="";
    length = recvfrom(extern_sock, buf, sizeof(buf), 0, (struct sockaddr *)&external, &length_extern);
    if (length > -1)
    {
        if (debug_level >= 1)
        {
            printf("\n\n---- Recv:Extern [IP:%s]----\n", inet_ntoa(external.sin_addr));
            time_t t = time(NULL);
            char temp[64];
            strftime(temp, sizeof(temp), "%Y/%m/%d %X %A", localtime(&t));
            printf("%s\n", temp);
            if (debug_level == 2)
                print_package(buf, length);
        }
        /* 根据返回结果修改请求记录表并且修改数据 */
        unsigned short ID;
        memcpy(&ID, buf, sizeof(unsigned short));
        client = id_trans_table[ID-1]->client;
        unsigned short true_ID = id_trans_table[ID-1]->old_ID;
        ID_Trans_Table *current = id_trans_table[ID-1];
        free(current);
        id_trans_table[ID-1] = NULL;
        memcpy(buf, &true_ID, sizeof(unsigned short));
        get_url_from_question(buf,url,debug_level);
        get_ip_from_answer(buf,ip,debug_level);
        if (strcmp(ip, "") != 0) {
            add_to_cache(cache, local_table, cache_lru_list, url, ip, debug_level);
            if (debug_level)
                printf("recieve packet [Url:%s -> IP:%s]\n", url, ip);
            length = sendto(local_sock, buf, length, 0, (SOCKADDR*)&client, sizeof(client));
            if (debug_level && length > 0)
                printf("Send packet to client[Url:%s -> IP:%s]\n", url, ip);
        }
        else {
            printf("The package doesn't have IPv4 address\n");
            length = sendto(local_sock, buf, length, 0, (SOCKADDR*)&client, sizeof(client));
        }
            
    }
    return;
}